package com.whz.governance.monitor;

import com.alibaba.csp.sentinel.slots.block.BlockException;
import com.google.common.collect.Lists;
import com.whz.generic.enhance.CommonClassEnhancer;
import javassist.CannotCompileException;
import javassist.CtClass;
import javassist.CtField;
import javassist.CtMethod;
import javassist.CtNewMethod;
import lombok.SneakyThrows;
import lombok.experimental.UtilityClass;
import lombok.extern.slf4j.Slf4j;
import lombok.val;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;

import java.lang.annotation.Annotation;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.util.Arrays;
import java.util.List;
import java.util.Optional;
import java.util.function.Predicate;
import java.util.stream.Collectors;

/**
 * 字节码增强器,创建sentinel的带BlockException的fallback方法
 *
 * @author whz
 */
@Slf4j
@val
@UtilityClass
public final class SentinelClassEnhancer {

    static final List<Method> OBJECT_METHODS = getObjectMethodList();

    private List<Method> getObjectMethodList() {
        try {
            return Lists.newArrayList(Object.class.getDeclaredMethods());
        } catch (Exception e) {
            log.error("", e);
        }
        return Lists.newArrayListWithCapacity(0);
    }

    /**
     * 增强现有的方法,增加它的fallback方法,sentinel额外增加 {@link BlockException} 参数
     *
     * @param clz         java class
     * @param annotations annotations on class
     * @param isSentinel  is sentinel filter
     * @return enhanced class
     */
    public static Class<?> fallbackMethodEnhance(final Class<?> clz,
                                                 List<Class<? extends Annotation>> annotations,
                                                 boolean isSentinel) {
        Assert.notNull(clz, "class cannot null");
        Assert.notEmpty(annotations, "annotations cannot empty");
        List<Method> lst;

        if (!isAnnotationPresent(annotations, clz)) {
            // 没有注解的,查看所有非静态public方法
            lst = getAllPublicMethods(clz);
        } else {
            lst = getAllPublicMethodsWithAnnotation(clz, annotations);
        }
        // 去掉object的方法
        lst.removeAll(OBJECT_METHODS);
        if (CollectionUtils.isEmpty(lst)) {
            return clz;
        }

        Optional<CtClass> option = safeObtainCtClass(clz.getName());
        if (!option.isPresent()) {
            return clz;
        }
        CtClass ctClass = option.get();
        // 对现有方法添加sentinel的异常参数
        lst.forEach(
                method -> addParamsEnhance(method, ctClass, isSentinel? BlockException.class : null));
        return CommonClassEnhancer.createJavaClass(clz.getName(), ctClass);
    }

    private static List<Method> getAllPublicMethodsWithAnnotation(Class<?> clz,
                                                                  List<Class<? extends Annotation>> annotations) {
        return Arrays.stream(clz.getDeclaredMethods()).filter(
                             method -> !Modifier.isStatic(method.getModifiers())
                                       && Modifier.isPublic(method.getModifiers())
                                       && isAnnotationPresent(annotations, method))
                     .collect(Collectors.toList());
    }

    private static List<Method> getAllPublicMethods(Class<?> clz) {
        return Arrays.stream(clz.getDeclaredMethods()).filter(
                             method -> !Modifier.isStatic(method.getModifiers())
                                       && Modifier.isPublic(method.getModifiers()))
                     .collect(Collectors.toList());
    }


    private static boolean isAnnotationPresent(List<Class<? extends Annotation>> annotations,
                                               Method method) {
        return isAnnotationPresent(annotations, method::isAnnotationPresent);
    }

    private static boolean isAnnotationPresent(List<Class<? extends Annotation>> annotations,
                                               Class<?> clz) {
        return isAnnotationPresent(annotations, clz::isAnnotationPresent);
    }

    private static boolean isAnnotationPresent(List<Class<? extends Annotation>> annotations,
                                               Predicate<Class<? extends Annotation>> predicate) {
        if (CollectionUtils.isEmpty(annotations)) {
            return false;
        }
        for (Class<? extends Annotation> annotation : annotations) {
            if (predicate.test(annotation)) {
                return true;
            }
        }
        return false;
    }

    /**
     * 对现有方法增强,添加参数
     */
    @SneakyThrows
    public static void addParamsEnhance(final Method method, CtClass ctClass, Class<?>... params) {
        try {
            CtMethod fMethod = ctClass.getDeclaredMethod(method.getName(), getParamsTypes(method));
            String newMethodName = createNewMethodName(method.getName());
            CtMethod gMethod = CtNewMethod.copy(fMethod, newMethodName, ctClass, null);
            if (params != null) {
                for (Class<?> p : params) {
                    Optional<CtClass> op = safeObtainCtClass(p.getName());
                    if (op.isPresent()) {
                        gMethod.addParameter(op.get());
                    }
                }
                // not static method
                int index = getParamIndex(method) + 1;
                gMethod.insertBefore("{ " +
                                     "throw new RuntimeException($" + index + ");" +
                                     " }");
            } else {
                gMethod.insertBefore("{ " +
                                     "throw new RuntimeException(\"flow rules protecting\");" +
                                     " }");
            }
            ctClass.addMethod(gMethod);
            ctClass.writeFile();
        } catch (Exception e) {
            log.error("add parameter to fall back method error.", e);
        }
    }

    private CtClass[] getParamsTypes(Method method) {
        CtClass[] parameterCtClasses = new CtClass[method.getParameterTypes().length];
        for (int i = 0; i < parameterCtClasses.length; i++) {
            parameterCtClasses[i] = CommonClassEnhancer.loadCtClass(
                    method.getParameterTypes()[i].getName());
        }
        return parameterCtClasses;
    }

    private int getParamIndex(Method method) {
        return method.getParameterCount();
    }

    /**
     * 服务降级方法名称统一为:'xxx_FallBack',xxx为原方法名称
     */
    private static String createNewMethodName(final String methodName) {
        return methodName + "TongyuFB";
    }

    public static Optional<CtClass> safeObtainCtClass(final String className) {
        try {
            return Optional.ofNullable(CommonClassEnhancer.loadCtClass(className));
        } catch (Exception e) {
            log.error("", e);
        }
        return Optional.empty();
    }

    /**
     * add field to ctClass
     *
     * @param ctClass   目标类
     * @param fieldName 新field的名称
     * @param fieldType 新field类型
     * @throws CannotCompileException 编译异常
     */
    public static void addField(CtClass ctClass, String fieldName, Class<?> fieldType)
            throws CannotCompileException {
        Optional<CtClass> fCtClass = safeObtainCtClass(fieldType.getName());
        if (!fCtClass.isPresent()) {
            return;
        }
        CtField newCtField = new CtField(fCtClass.get(), fieldName, ctClass);
        newCtField.setModifiers(Modifier.PRIVATE);
        ctClass.addField(newCtField);
    }
}